# location based representation of object
import random
import uuid
from collections import defaultdict
import torch

import numpy as np
from skimage import measure

import visionAndTouch.utils as utils
import copy

class CauseEffectPosition:

    def __init__(self, causeEffectId, actionId, objInstId, positionDelta):
        self.id = causeEffectId
        self.actionId = actionId
        self.objInstId = objInstId
        self.positionDelta = positionDelta

    def __eq__(self, other):
        if self.actionId != other.actionId:
            return False
        # if self


class ObjectDescription:

    def __init__(self, objDescriptionId, coordinates, feature):
        self.id = objDescriptionId
        self.locations = set()
        self.feature = feature
        self.coordinates = coordinates
        self.locationToCoordinateMap = defaultdict(object)
        for coord in coordinates:
            loc = uuid.uuid1()
            self.locations.add(loc)
            self.locationToCoordinateMap[loc] = (coord[0:2])


class ObjectInstance:

    def __init__(self, objInstanceId, objDescriptionId, bbox):
        self.id = objInstanceId
        self.objDescriptionId = objDescriptionId
        self.bbox = bbox

    def getBboxTopLeft(self):
        return self.bbox[0:2]


class BodyInfo:

    def __init__(self, regionInfo):
        self.area = regionInfo.area
        self.area_bbox = regionInfo.area_bbox


class MemoryInfo:

    def __init__(self):
        pass


class SpaceInfo:

    def __init__(self, position, velocity):
        self.position = position
        self.velocity = velocity


class ObservationInfo:

    def __init__(self):
        self.beforeSpaceInfo = None
        self.action = None
        self.afterSpaceInfo = None


class Agent:

    def __init__(self, worldDim):
        self.worldHeight = worldDim[0]
        self.worldWidth = worldDim[1]
        self.learnedObjectDescriptions = []  # need to be persisted
        self.matchedObjInst = []
        self.maxODID = 0
        self.maxObjInstId = 0
        self.maxCauseEffectId = 0
        self.currentScene = []
        self.previousScene = []
        self.actionSpace = [0, 1, 2, 3]
        self.previousAction = None
        self.learnedCauseEffects = []
        self.temporalMemory = []
        self.constructedObjects = []
        self.previousScene = None
        self.currentScene = None
        self.memoryInfo = None
        self.bodyInfo = None
        self.internalTime = 0
        self.predictionRate = 0
        self.previousSpaceInfo = None
        self.currentSpaceInfo = None
        self.observations = []
        # self.isBodyStatic = False
        # self.isAct0ExperimentDone = False
        # self.act0ExperimentData = []
        # self.act0ExperimentStartIter = None
        # self.act0ExperimentEndIter = None

    def detectMotionBetweenTwoScene(self):
        movedRegions = []
        for cr in self.currentScene:
            pr = Agent.findRegionByArea(cr.area, self.previousScene)
            if pr is not None:  # found an regions with the same area
                cbbox = cr.bbox
                pbbox = pr.bbox
                if np.sum(np.array(pbbox) - np.array(cbbox)) != 0:
                    movedRegions.append(cr)
        return movedRegions

    @staticmethod
    def findRegionByArea(area, regions):
        for r in regions:
            if area == r.area:
                return r
        return None

    @staticmethod
    def findRegionByBbox(bbox, regions):
        for r in regions:
            if bbox == r.bbox:
                return r
        return None

    def handle(self, observation):
        self.internalTime += 1
        scene = Agent.getRegionsInTheScene(observation)
        if self.previousScene is None:
            self.previousScene = scene
            return
        self.currentScene = scene
        # find your body first: in the scene, only your body can move
        while self.bodyInfo is None:
            movedRegions = self.detectMotionBetweenTwoScene()
            if len(movedRegions) == 0 or len(movedRegions) > 1:  # if no motion, make some noise
                return self.chooseRandomAction()
            elif len(movedRegions) == 1:  # found motions, mark the moved region as body
                self.bodyInfo = BodyInfo(movedRegions[0])
                torch.save(self.bodyInfo, 'savedModels/bodyInfo.m')
                break
        while self.predictionRate < 9.0:
            bodyInTheScene = self.findBodyInTheScene(scene)
            currentPosition = bodyInTheScene.centroid
            if len(self.observations) == 0:  # the first observation
                act = self.chooseRandomAction()
                obInfo = ObservationInfo()
                obInfo.beforeSpaceInfo = SpaceInfo(currentPosition, None)
                obInfo.action = act
                self.observations.append(obInfo)
                return act
            else:
                # finish the last observation and start a new one
                lastObservation = self.observations[-1]
                previousPosition = lastObservation.beforeSpaceInfo.position
                velocity = np.array(currentPosition) - np.array(previousPosition)
                lastObservation.afterSpaceInfo = SpaceInfo(currentPosition, velocity)
                # start a new observation
                act = self.chooseRandomAction()
                newObInfo = ObservationInfo()
                newObInfo.beforeSpaceInfo = copy.deepcopy(lastObservation.afterSpaceInfo)
                newObInfo.action = act
                self.observations.append(newObInfo)
                return act
        return
        # once you found your body, you need to find your body in every scene. your body is your anchor.
        # accumulate temporal memory
        # while not self.isBodyStatic:
        #     self.temporalMemory.append(scene)
        #     if len(self.temporalMemory) > 10:
        #         self.temporalMemory.pop(0)
        #     if len(self.temporalMemory) < 10:
        #         return
        #     bodyInEachScene = self.findBodyBboxInTheTemporalMemory(self.temporalMemory)
        #     if len(bodyInEachScene) == 1:
        #         self.isBodyStatic = True
        #         break
        #     else:
        #         return
        # while not self.isAct0ExperimentDone:
        #     if self.act0ExperimentStartIter is None:
        #         self.act0ExperimentStartIter = self.internalTime + 1
        #         return 0
        #     if self.internalTime >= self.act0ExperimentStartIter:
        #         self.act0ExperimentData.append(scene)
        #     return

    def findBodyBboxInTheTemporalMemory(self, temporalMemory):
        result = set()
        for tm in temporalMemory:
            result.add(tuple(self.findBodyInTheScene(tm).bbox))
        return result

    def findBodyInTheScene(self, theScene):
        for r in theScene:
            if r.area == self.bodyInfo.area:
                return r
        for r in theScene:
            if r.area_bbox == self.bodyInfo.area_bbox:
                return r
        # can not find body, print the reason:
        for r in theScene:
            print('area_bbox', r.area_bbox)
        raise Exception('cann not find body')

    def chooseRandomAction(self):
        action = random.randint(0, len(self.actionSpace) - 1)
        return action

    @staticmethod
    def getRegionsInTheScene(theScene):
        encodedImageT = utils.encodeImageTorch(theScene)
        encodedImage = encodedImageT.cpu().numpy()
        label = measure.label(encodedImage, connectivity=2)
        props = measure.regionprops(label)
        return props

    def createCurrentScene(self, theScene):
        self.previousScene = self.currentScene
        self.currentScene = []
        encodedImageT = utils.encodeImageTorch(theScene).unsqueeze(dim=2)
        encodedImage = encodedImageT.numpy()
        label = measure.label(encodedImage, connectivity=2)
        props = measure.regionprops(label)
        # matching regions with learned objects
        for pp in props:
            firstCoord = pp.coords[0][0:2]
            feature = theScene[firstCoord[0], firstCoord[1]]
            if len(self.learnedObjectDescriptions) == 0:
                # if there is no learned object description, just create a new one
                objDescrip = ObjectDescription(self.maxODID, pp.coords, feature)
                self.maxODID += 1
                self.learnedObjectDescriptions.append(objDescrip)
                # after create a new lod, create a new object instance
                objInst = ObjectInstance(self.maxObjInstId, objDescrip.id, pp.bbox)
                self.maxObjInstId += 1
                self.currentScene.append(objInst)
                continue
            foundMatchForTheRegion = False
            for objDescrip in self.learnedObjectDescriptions:
                # if match, create an object instance
                if self.matchDescription(pp.coords, feature, objDescrip):
                    objInst = ObjectInstance(self.maxODID, objDescrip.id, pp.bbox)
                    self.maxODID += 1
                    self.currentScene.append(objInst)
                    foundMatchForTheRegion = True
                    break
            if not foundMatchForTheRegion:
                # if do not match, create a new learned object description
                objDescrip = ObjectDescription(self.maxODID, pp.coords, feature)
                self.maxODID += 1
                self.learnedObjectDescriptions.append(objDescrip)
                objInst = ObjectInstance(self.maxObjInstId, objDescrip.id, pp.bbox)
                self.maxObjInstId += 1
                self.currentScene.append(objInst)

    def matchObjectBetweenScene(self):
        for objInstPrevious in self.previousScene:
            matchingObjInst = Agent.findNoChangesInTheScene(objInstPrevious, self.currentScene)
            if matchingObjInst:
                matchingObjInst.id = objInstPrevious.id
            else:  # not matching object that stay in the old place, check position change:
                positionChangedObjInst = Agent.findPositionChanges(objInstPrevious, self.currentScene)
                positionChangedObjInst.id = objInstPrevious.id

    def learnCauseEffect(self):
        if self.previousAction is None or len(self.previousScene) == 0:
            return
        for currenObjInst in self.currentScene:
            previousObjInst = self.findObjInstById(currenObjInst.id, self.previousScene)
            currentBbox = currenObjInst.bbox
            previousBbox = previousObjInst.bbox
            if currentBbox != previousBbox:
                existingCEP = self.getCauseEffectPosition(self.previousAction)
                if not existingCEP:
                    positionDelta = (np.array(currentBbox) - np.array(previousBbox))[0:2]
                    newcep = CauseEffectPosition(self.maxCauseEffectId, self.previousAction, currenObjInst.id, positionDelta)
                    self.maxCauseEffectId += 1
                    self.learnedCauseEffects.append(newcep)

    def getCauseEffectPosition(self, actionId):
        for cep in self.learnedCauseEffects:
            if cep.actionId == actionId:  # each atomic action has only one causal rule
                return cep
        return None

    @staticmethod
    def findObjInstById(objInstId, objInstList):
        for objInst in objInstList:
            if objInst.id == objInstId:
                return objInst

    @staticmethod
    def findNoChangesInTheScene(objInst, theScene):
        for objInstInTheScene in theScene:
            if objInst.objDescriptionId == objInstInTheScene.objDescriptionId and objInst.bbox == objInstInTheScene.bbox:
                return objInstInTheScene
        return None

    @staticmethod
    def findPositionChanges(objInst, theScene):
        for objInstInTheScene in theScene:
            if objInst.objDescriptionId == objInstInTheScene.objDescriptionId and objInst.bbox != objInstInTheScene.bbox:
                return objInstInTheScene
        return None

    def matchDescription(self, ppCoord, ppFeature, objectDescription):
        if utils.encodeRGB(objectDescription.feature) != utils.encodeRGB(ppFeature):
            return False
        if ppCoord.shape[0] != objectDescription.coordinates.shape[0]:
            return False
        displacement = ppCoord - objectDescription.coordinates
        firstDiff = utils.encodeCoord(displacement[0], self.worldWidth)
        for d in displacement:
            dEncoded = utils.encodeCoord(d, self.worldWidth)
            if firstDiff != dEncoded:
                return False
        return True

    def getAction(self, targetObjInstId, targetBboxTopLeft):
        if len(self.learnedCauseEffects) < len(self.actionSpace):
            return self.chooseRandomAction()
        else:
            return self.plan(targetObjInstId, targetBboxTopLeft)

    def plan(self, targetObjInstId, targetBboxTopLeft):
        targetObjInst = self.findObjInstById(targetObjInstId, self.currentScene)
        currBboxTopLeft = targetObjInst.getBboxTopLeft()
        currManhattan = utils.manhattan(targetBboxTopLeft, currBboxTopLeft)
        availableAction = []
        # strategy one: try out actions in the mind
        for act in self.actionSpace:
            cep = self.findCauseEffectByActionId(act)
            updatedBboxTopLeft = currBboxTopLeft + cep.positionDelta
            updatedManhattan = utils.manhattan(targetBboxTopLeft, updatedBboxTopLeft)
            if updatedManhattan < currManhattan:
                availableAction.append(act)
        selectAct = random.randint(0, len(availableAction)-1)
        return availableAction[selectAct]

    def checkTargetFininsed(self, targetObjInstId, targetBboxTopLeft):
        targetObjInst = self.findObjInstById(targetObjInstId, self.currentScene)
        currBboxTopLeft = targetObjInst.getBboxTopLeft()
        currManhattan = utils.manhattan(targetBboxTopLeft, currBboxTopLeft)
        if currManhattan == 0:
            print('reached target, no action')
            return True
        else:
            return False

    def findCauseEffectByActionId(self, actionId):
        for cep in self.learnedCauseEffects:
            if cep.actionId == actionId:
                return cep

    @staticmethod
    def printScene(theScene):
        print('there are', len(theScene), 'objects in the scene')
        for ith, o in enumerate(theScene):
            print(str(ith) + ' object instance Id is ', o.id, '; object description Id is', o.objDescriptionId)
